import torch
from torch import nn

class RNN(nn.Module):
    def __init__(self, input_dim=128, hidden_dim=64, num_layers=1):
        super(RNN, self).__init__()
        self.rnn = nn.RNN(input_dim, hidden_dim, num_layers, batch_first=True)

    def forward(self, feature, mask):
        """
        Args:
            feature (Tensor): The input features, of shape [B, L, input_dim].
            mask (Tensor): The mask tensor, of shape [B, L].
            
        Returns:
            Tensor: The output of the RNN, of shape [B, L, hidden_dim].
        """
        # Apply mask to the input feature
        feature = feature * mask.unsqueeze(-1)
        output, _ = self.rnn(feature)
        return output

if __name__ == "__main__":
    # Test the RNN module
    input_dim = 128
    hidden_dim = 64
    model = RNN(input_dim, hidden_dim)

    B, L, dim = 32, 100, 128  
    feature = torch.randn(B, L, dim)
    mask = torch.ones(B, L)

    output = model(feature, mask)
    print(output.shape)